package cn.xingyun.pickaxe_loli_new.mixin;

import cn.xingyun.pickaxe_loli_new.item.LolipickaxeItem;
import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
import net.minecraft.client.Minecraft;
import net.minecraft.client.multiplayer.ClientLevel;
import net.minecraft.server.level.ServerLevel;
import net.minecraft.world.entity.Entity;
import net.minecraft.world.entity.player.Player;
import net.minecraft.world.level.entity.EntityTickList;
import net.minecraft.world.level.entity.LevelEntityGetter;
import net.minecraft.world.level.entity.PersistentEntitySectionManager;
import net.minecraft.world.level.entity.TransientEntitySectionManager;
import net.minecraftforge.entity.PartEntity;
import org.spongepowered.asm.mixin.Mixin;
import org.spongepowered.asm.mixin.Shadow;
import org.spongepowered.asm.mixin.injection.At;
import org.spongepowered.asm.mixin.injection.Inject;
import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;

public class Tick {
    @Mixin(value = ClientLevel.class)
    public static abstract class ClientLevelMixin {
        @Shadow protected abstract LevelEntityGetter<Entity> getEntities();

        @Shadow
        public TransientEntitySectionManager<Entity> entityStorage;

        @Shadow
        public Int2ObjectMap<PartEntity<?>> partEntities;

        @Shadow
        public EntityTickList tickingEntities;

        @Inject(method = "tick",at = @At("HEAD"))
        public void tick(CallbackInfo ci) {
            if (LolipickaxeItem.strongerRemove) {
                for (Entity entity : getEntities().getAll()) {
                    if (entity == null) {
                        return;
                    }
                    Player player = Minecraft.getInstance().player;
                    entityStorage.entityStorage.byId.clear();
                    entityStorage.entityStorage.byUuid.clear();
                    if (player != null) {
                        entityStorage.entityStorage.byId.put(player.getId(), player);
                        entityStorage.entityStorage.byUuid.put(player.getUUID(), player);
                    }
                    entityStorage.callbacks.onTickingEnd(entity);
                    entityStorage.callbacks.onTrackingEnd(entity);
                    entityStorage.callbacks.onDestroyed(entity);
                    partEntities.clear();
                    tickingEntities.remove(entity);
                }
            }
        }
    }
    @Mixin(value = ServerLevel.class)
    public static abstract class ServerLevelMixin{
        @Shadow public abstract Iterable<Entity> getAllEntities();

        @Shadow public PersistentEntitySectionManager<Entity> entityManager;

        @Shadow public EntityTickList entityTickList;

        @Inject(method = "tick",at = @At("HEAD"))
        public void tick(CallbackInfo ci) {
            if (LolipickaxeItem.strongerRemove) {
                for (Entity entity : getAllEntities()) {
                    if (entity == null) {
                        return;
                    }
                    Player player = Minecraft.getInstance().player;
                    entityManager.visibleEntityStorage.byId.clear();
                    entityManager.visibleEntityStorage.byUuid.clear();
                    if (player != null) {
                        entityManager.visibleEntityStorage.byId.put(player.getId(), player);
                        entityManager.visibleEntityStorage.byUuid.put(player.getUUID(), player);
                    }
                    entityManager.callbacks.onTickingEnd(entity);
                    entityManager.callbacks.onTrackingEnd(entity);
                    entityManager.callbacks.onDestroyed(entity);
                    entityTickList.remove(entity);
                }
            }
        }
    }
}
